import concurrent.futures
import atexit
import importlib
import os
import sys
import time
import json
import matplotlib.pyplot as plt
import matplotlib
import warnings

warnings.filterwarnings('ignore', category=FutureWarning)
#matplotlib.use('TkAgg')

from a1_planner import Planner as Planner_a1
from a2_planner import Planner as Planner_a2
from lqr_control import LQRController
from a1_simple_control import MotionController as MotionController_a1
from a2_simple_control import MotionController as MotionController_a2
from predictor import Predictor

def add_path_to_sys(target: str):
    abs_path = os.path.abspath(target)
    if abs_path not in sys.path:
        sys.path.append(abs_path)


def check_dir(target_dir):
    """check path"""
    if not os.path.exists(target_dir):
        os.makedirs(target_dir, exist_ok=True)


def exit_handler(executor):
    print("Exiting program...")
    executor.shutdown(wait=False)


def test_scenario(scenario_to_test, dir_save_img, dir_parent_1):
    try:
        # 在子进程中初始化模块和对象
        scenarioOrganizer = importlib.import_module("onsite-mine.dynamic_scenes.scenarioOrganizer1")
        env = importlib.import_module("onsite-mine.dynamic_scenes.env")
        lookup = importlib.import_module("onsite-mine.dynamic_scenes.lookup")
        envi = env.Env()

        so = scenarioOrganizer.ScenarioOrganizer()
        dir_inputs = os.path.abspath(os.path.join(dir_parent_1, 'inputs'))
        dir_outputs = os.path.abspath(os.path.join(dir_parent_1, 'outputs'))
        so.load(dir_inputs, dir_outputs)

        print(f"###log### <scene-{scenario_to_test['data']['scene_name']}>\n")
        formatted_so_config = json.dumps(so.config, indent=4, ensure_ascii=False)
        print(f"###log### <测试参数>\n{formatted_so_config}\n")
        collision_lookup = lookup.CollisionLookup()
        observation, traj, client = envi.make(scenario=scenario_to_test, save_img_path=dir_save_img,
                                              kinetics_mode='simple')
        if observation['test_setting']['scenario_type'] == 'intersection':
            ########### 算法 部分1:初始化 ###########
            # 交叉路口单车通行任务:构建整个场景的预测器,预测所有的车辆;构建规划器\控制器,控制ego车
            predictor = Predictor(time_horizon=5.0)
            planner = Planner_a1(observation)
            path_planned, spd_planned = planner.process(observation)
            lat_controller = MotionController_a1()
            step_sum = int(observation["test_setting"]['max_t'] / observation["test_setting"]['dt'])
            ########### 算法 部分1:初始化 ###########

            # 逐帧进行仿真，触发仿真停止条件时结束
            # 当测试还未进行完毕,即观察值中test_setting['end']还是-1的时候
            while observation['test_setting']['end'] == -1:
                ########### 算法 部分2:执行 ###########
                traj_future, _ = predictor.predict(observation, traj, predictor_type="CACV_PREDICTOR")
                action = lat_controller.process(observation['vehicle_info'], path_planned, traj_future, observation,
                                                traj)
                # action = (3,-0.55)
                action = action + (1,)  # 添加档位信息，默认为前进档位
                ########### 算法 部分2:执行 ###########
                observation = envi.step(action, traj_future, observation, traj)  # 根据车辆的action,更新场景,并返回新的观测值.
        else:
            Reverse = False
            simple = False
            planner = Planner_a2(observation)
            path_planned, spd_planned = planner.process(collision_lookup, observation,traj=traj)

            scene_name = scenario_to_test['data']['scene_name']
            task = scene_name[scene_name.find("_") + 1:scene_name.find("_", scene_name.find("_") + 1)]
            ind_list = []
            while observation['test_setting']['end'] == -1:
                if task == "loading":
                    if Reverse == False:
                        controller = LQRController()
                        action, Reverse, ind_list = controller.process(observation['vehicle_info'], path_planned,
                                                                       spd_planned,
                                                                       "simple", Reverse=Reverse, simple=simple,
                                                                       ind_list=ind_list)
                    else:
                        controller = MotionController_a2(20)
                        action, ind, Reverse, arrive = controller.process(observation['vehicle_info'],
                                                                          path_planned,
                                                                          spd_planned, observation,
                                                                          Reverse=Reverse,
                                                                          traj=traj,
                                                                          scenario_to_test=scenario_to_test)
                        if arrive:
                            Reverse = False
                            simple = True
                            print("切换控制方法为LQR")
                elif task == "unloading":
                    controller = MotionController_a2(10)
                    action, _, _, _ = controller.process(observation['vehicle_info'], path_planned,
                                                         spd_planned, observation, Reverse=False, traj=traj,
                                                         scenario_to_test=scenario_to_test)
                observation = envi.step(action, 100, observation, 100)

    except Exception as e:
        print(repr(e))
        if client is not None:
            client.close_sockets()
    finally:
        so.add_result(scenario_to_test, observation['test_setting']['end'])
        plt.close()


def main():
    dir_current_file = os.path.dirname(__file__)
    dir_parent_1 = os.path.dirname(dir_current_file)
    dir_save_img = os.path.abspath(os.path.join(dir_parent_1, 'onsite_images_saved'))

    tic = time.time()

    # 加载场景
    scenarioOrganizer = importlib.import_module("onsite-mine.dynamic_scenes.scenarioOrganizer1")
    so = scenarioOrganizer.ScenarioOrganizer()
    dir_inputs = os.path.abspath(os.path.join(dir_parent_1, 'inputs'))
    dir_outputs = os.path.abspath(os.path.join(dir_parent_1, 'outputs'))
    so.load(dir_inputs, dir_outputs)

    # 准备测试场景列表
    scenario_list = []
    while True:
        scenario_to_test = so.next()
        if scenario_to_test is None:
            break
        scenario_list.append(scenario_to_test)

    executor = concurrent.futures.ProcessPoolExecutor(max_workers=40)
    atexit.register(exit_handler, executor)
    with executor as executor_inner:
        futures = []
        for scenario_to_test in scenario_list:
            futures.append(executor_inner.submit(test_scenario, scenario_to_test, dir_save_img, dir_parent_1))
        for future in concurrent.futures.as_completed(futures):
            pass
    toc = time.time()
    print(f"###log### 总用时: {int((toc - tic) // 3600)}小时 {int(((toc - tic) % 3600) // 60)}分钟 {(toc - tic) % 60:.2f}秒\n")


if __name__ == "__main__":
    main()
